'''
1、安装模块
pip install baidu-aip
pip install opencv-python
2、申请百度开放平台,创建应用
https://ai.baidu.com/ai-doc/OCR/dk3iqnq51
3、替换自己的APP_ID、API_KEY、SECRET_KEY
'''
from aip import AipOcr
import os
import shutil
import cv2
import numpy as np
import ctypes

from util import ReadJsonFile

# 可选参数(忽略即可)
options = {}
# options["recognize_granularity"] = "big"
# options["language_type"] = "CHN_ENG"
# options["detect_direction"] = "true"
# options["detect_language"] = "true"
options["vertexes_location"] = "true"
# options["probability"] = "true"

# 读取图片所有内容
def __get_file_content(filePath):
  with open(filePath, 'rb') as fp:
    return fp.read()

# 拼接2张图片
def __connact_image(img1, img2):
  height1, width1, _ = img1.shape
  height2, width2, _ = img2.shape

  # new image
  final_matrix = np.zeros((height1, width1+width2, 3), np.uint8)
  # change 
  final_matrix[0:height1, 0:width1] = img1
  final_matrix[0:height1, width1:width1+width2] = img2
  return final_matrix

# 指定位置, 裁剪图片, 保存为一个中间文件
def __crop_image(img, x, y, width, height, n, prefix_name, offset, gray, mask_img=None):
  image_files = []
  x0, y0 = x, y
  left, right, top, bottom = offset
  try:
    for i in range(n):
      x1 = x0 + width
      y1 = y0 + height
      cell_img_file = '%s-%d.png' % (prefix_name, i)
      # FIX 截取范围向右偏1个像素，下偏1个像素
      cell_img = img[y0+top:y1+bottom, x0+left:x1+right]
      # 放大2倍
      # cell_img = cv2.resize(cell_img,None,fx=1,fy=1)
      # 灰度化、二值化处理
      if gray:
        # 图像灰度化(降低图片像素大小)
        cell_img = cv2.cvtColor(cell_img, cv2.COLOR_RGB2GRAY)
        # 直接阈值化是对输入的单通道矩阵逐像素进行阈值分割(效果较差!!!)
        # ret, cell_img = cv2.threshold(cell_img, 50, 255, cv2.THRESH_BINARY | cv2.THRESH_TRIANGLE)
        # 自适应阈值化能够根据图像不同区域亮度分布，改变阈值(效果较差!!!)
        # cell_img =  cv2.adaptiveThreshold(cell_img, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2)
      # 如果mask_img不为空,则表示使用mask增强算法, mask图片和原图进行拼接, 提高识别率
      if mask_img is not None:        
        cell_img = __connact_image(mask_img, cell_img)
      # 保存图片
      cv2.imencode('.png', cell_img)[1].tofile(cell_img_file)
      # 保存图片(缺点: 不可以保存中文路径)
      # cv2.imwrite(cell_img_file, img[y0:y1, x0:x1])
      image_files.append(cell_img_file)
      y0 = y1
  except Exception as err:
    print('[错误]crop_image:', err)
    image_files = []
  return image_files

# 获取单元格的位置和尺寸
def __get_cell_rect(config, key):
  return config[key]['pos']['x'], config[key]['pos']['y'], config['cell']['width'], config['cell']['height']

# 构造单元格截图的文件名
def __build_crop_image_filename(name, key, config):
  return '%s\\%s-%s' % (config['crop_dir'], name, key)

def crop_cell_images(image_path, filename, config):
  image_files = []
  # opencv读取图片
  img = cv2.imdecode(np.fromfile(image_path, dtype=np.uint8), 1)
  if config['mask_enhance']:
    # opencv读取要拼接图片
    mask_img = cv2.imdecode(np.fromfile(config['mask_image'], dtype=np.uint8), 1)
  else:
    mask_img = None

  # 截图xyz坐标的每一个单元格
  x, y, width, height = __get_cell_rect(config, 'xyz')
  crop_image_filename = __build_crop_image_filename(filename, 'xyz', config)
  image_files.extend(__crop_image(img, x, y, width, height, len(config['xyz']['titles']), crop_image_filename, config['cell']['offset'], config['gray'], mask_img))

  # 截图巷道数据的每一个单元格
  x, y, width, height = __get_cell_rect(config, 'tunnel')
  crop_image_filename = __build_crop_image_filename(filename, 'tunnel', config)
  image_files.extend(__crop_image(img, x, y, width, height, len(config['tunnel']['titles']), crop_image_filename, config['cell']['offset'], config['gray'], mask_img))
  # 返回
  return image_files

# 图片第1个字符是否有负号
def image_has_negative(cell_img_file, first_char_width, max_black_pixels):
  # opencv读取图片
  img = cv2.imdecode(np.fromfile(cell_img_file, dtype=np.uint8), 1)
  # 截取第一个字符
  # height, weight, channels = img.shape
  pixel_data = np.array(img[0:img.shape[0], 0:first_char_width], dtype = np.uint8)
  # print(pixel_data)
  # 黑色像素的个数(使用numpy的sum函数统计)
  black_pixel_nums = np.sum(pixel_data<255)
  # print(img.shape, '截图的像素和:', zero_nums)
  # 小于30个,则认为图片中的第一个是负号
  return black_pixel_nums < max_black_pixels

def get_result(word):
  # 去掉字符串中的空格
  word.replace(' ', '')
  # 判断识别结果是否以'NULL-'开头
  pos = word.find('#')
  # FIX 有些情况NULL可能会被识别为NUL(2个L靠的太近被认为是1个L)
  # and (word[0:pos]=='NULL' or word[0:pos]=='NUL')
  if pos > -1:
    # 取NULL-之后的数据作为识别结果
    return word[pos+1:].strip()
  else:
    return word.strip()

# 利用百度AI文字识别API提取图片数据
def ocr_images(image_files, client):
  data = []
  for cell_img_file in image_files:
    # print(cell_img_file)
    try:
      # 文字识别
      message = client.basicGeneral(__get_file_content(cell_img_file), options)
      # 解决Open api qps request limit reached问题
      while True:
        if 'error_code' in message:
          if message['error_code'] == 18:
            message = client.basicGeneral(__get_file_content(cell_img_file), options)
          else:
            break
        else:
          break
      # 取出识别结果
      # 初步判断,其它的错误交给异常处理
      if 'words_result' not in message:
        data.append('NULL')
      elif len(message['words_result']) == 0:
        data.append('NULL')
      elif 'words' not in message['words_result'][0]:
        data.append('NULL')
      else:
        word = message['words_result'][0]['words']
        # DEPRECATED 判断负号的方法已废弃!
        # if word[0] != '-' and image_has_negative(cell_img_file, first_char_width, max_black_pixels):
        #   print(cell_img_file+' 负号未识别, 自动补上!')
        #   # 将负号补回去
        #   word = '%s%s' % ('-', word)
        # 提取结果
        # 说明: word理论都以NULL-或者NUL-开头!!!
        new_word = get_result(word)
        # 添加到数组
        data.append(new_word)
        print(cell_img_file, '--> 处理前:', word, ' 处理后:', new_word)
    except Exception as err:
      print('[错误]ocr_image:', err)
      data.append('NULL')
  return data

def get_datas_from_image(client, image_path, filename, config):
  # 截取所有单元格的图片
  image_files = crop_cell_images(image_path, filename, config)
  # 判断单元格总数与图片个数是否一致
  n1, n2 = len(config['xyz']['titles']), len(config['tunnel']['titles'])
  if n1+n2 != len(image_files):
    return None
  if config['ocr']:
    # 从图片中提取单元格数据
    return ocr_images(image_files, client)
  else:
    return None

# 提取数据标题
def get_titles(config):
  titles = []
  titles.extend(config['xyz']['titles'])
  titles.extend(config['tunnel']['titles'])
  return titles

# 主函数
def main(config):
  # 删除并新建截图文件夹
  shutil.rmtree(config['crop_dir'], ignore_errors=True)    #递归删除文件夹
  if not os.path.exists(config['crop_dir']):
    # os.rmdir(config['crop_dir'])
    os.mkdir(config['crop_dir'])

  # 创建百度AI文字识别客户端
  client = AipOcr(config['APP_ID'], config['API_KEY'], config['SECRET_KEY'])

  # 属性数据二维数组
  props = []
  # 第一行: 标题
  props.append(get_titles(config))

  # 遍历图片文件夹
  for parent, dirnames, filenames in os.walk(config['image_dir'], followlinks=True):
    nCount = 1
    for _filename in filenames:
      # 判断扩展名类型是否图片
      filename, ext = os.path.splitext(_filename)
      if ext != '.png' and ext != '.jpg' and ext != '.bmp':continue

      # 完整路径
      image_path = os.path.join(parent, _filename)
      print('%d: 处理图片 \"%s\"' % (nCount, image_path))

      # 从图片提取数据
      filename = str(nCount)
      props.append(get_datas_from_image(client, image_path, filename, config))

      # 计数器
      nCount += 1
  # 保存数据文件
  np.savetxt(config['out_file'], props, delimiter='\t', fmt='%s')

if __name__ == '__main__':
  # 配置参数
  config = ReadJsonFile('ocr_config.json')
  # 主函数
  main(config)
  # 暂停
  os.system("pause")